from abc import ABC, abstractmethod
import os
import pathlib
from typing import Union

from pyschism.server import ServerConfig, SlurmConfig


class Makefile(ABC):

    def __init__(self, server_config: ServerConfig = None, hotstart=None):
        if server_config is None:
            server_config = ServerConfig()
        self.server_config = server_config
        self.hotstart = hotstart

    def write(self, path: Union[str, os.PathLike], overwrite: bool = False):
        path = pathlib.Path(path)
        if path.exists() and overwrite is not True:
            raise IOError(
                f"File {str(path)} exists and overwrite is not True.")
        if self.hotstart is not None:
            self._relpath = pathlib.Path(os.path.relpath(self.hotstart.path.parent, path.parent))
        with open(path, 'w') as f:
            f.write(str(self))
        if hasattr(self, '_relpath'):
            del self._relpath

    @property
    @abstractmethod
    def run(self):
        """Makefile run target."""

    @property
    def tail(self):
        return """
tail:
    tail -f outputs/mirror.out  outputs/fatal.error
"""

    @property
    def symlinks(self):
        return r"""
symlinks:
    @set -e;\
    if [ ! -z $${SYMLINK_OUTPUTS_DIR} ];\
    then \
        ln -sf $${SYMLINK_OUTPUTS_DIR} $${ROOT_DIR}outputs;\
    else \
        mkdir -p $${ROOT_DIR}outputs;\
    fi;\
    touch outputs/mirror.out outputs/fatal.error
"""


class DefaultMakefile(Makefile):

    def __str__(self):
        f = [
            "# Makefile driver generated by PySCHISM.",
            r"MAKEFILE_PATH:=$(abspath $(lastword $(MAKEFILE_LIST)))",
            r"ROOT_DIR:=$(dir $(MAKEFILE_PATH))",
            str(self.server_config),
            # self.default,
            # self.symlinks,
            self.run,
            self.tail,
        ]
        return "\n".join([line.replace("    ", "\t") for line in f])

#     @property
#     def default(self):
#         return r"""
# default: symlinks
# """

    @property
    def run(self):
        f = [
            '',
            'run:',
            '    @set -e;\\',
        ]
        if self.server_config.modules_init is not None:
            f.append(f'    source {self.server_config.modules_init};\\')
        if self.server_config.modulepath is not None:
            f.append(f'    export MODULEPATH={self.server_config.modulepath};\\')
        if self.server_config.modules is not None:
            f.append(f'    module load {" ".join(module for module in self.server_config.modules)};\\')
        if self.hotstart is not None:
            f.extend([
                f'    pushd {self._relpath if hasattr(self, "_relpath") else self.hotstart.path.parent.resolve()};\\',
                f'    {self.hotstart.binary} -i {self.hotstart.iteration};\\',
                '    popd;\\',
                f'    mv {self._relpath / self.hotstart.path.name if hasattr(self, "_relpath") else self.hotstart.path.resolve()} ./hotstart.nc;\\',
            ])

        return '\n'.join([line.replace("    ", "\t") for line in f]) + r"""
    rm -rf outputs/mirror.out outputs/fatal.error;\
    touch outputs/mirror.out outputs/fatal.error;\
    eval 'tail -f outputs/mirror.out  outputs/fatal.error &';\
    tail_pid=$${!};\
    ${MPI_LAUNCHER} ${NPROC} ${SCHISM_BINARY};\
    err_code=$${?};\
    kill "$${tail_pid}";\
    exit $${err_code}
"""


class SlurmMakefile(Makefile):

    def __str__(self):
        f = [
            "# Makefile driver generated by PySCHISM.",
            r"MAKEFILE_PATH:=$(abspath $(lastword $(MAKEFILE_LIST)))",
            r"ROOT_DIR:=$(dir $(MAKEFILE_PATH))",
            str(self.server_config),
            '',
            self.default,
            '',
            # self.symlinks,
            self.slurm,
            self.run,
            self.tail,
        ]
        return "\n".join([line.replace("    ", "\t") for line in f])

    @property
    def default(self):
        return 'default: run'

    @property
    def slurm(self):
        def indent(n, string):
            return n*' ' + string

        def modules():
            f = []
            if self.server_config.modules_init is not None:
                f.append(indent(4, fr'printf ". {self.server_config.modules_init}\n" >> ' + '${SLURM_JOB_FILE};' + '\\'))
            if self.server_config.modulepath is not None:
                f.append(indent(4, fr'printf "export MODULEPATH={self.server_config.modulepath}\n" >> ' + r'${SLURM_JOB_FILE};' + '\\'))
            if self.server_config.modules is not None:
                f.append(indent(4, 'printf "module load ' + fr'{" ".join([module for module in self.server_config.modules])}\n" >> ' + r'${SLURM_JOB_FILE};' + '\\'))
            return f

        def hotstart():
            f = []
            if self.hotstart is not None:
                return [
                    indent(4, fr'printf "pushd {self._relpath if hasattr(self, "_relpath") else self.hotstart.path.parent.resolve()}\n" >> ' + '${SLURM_JOB_FILE};' + '\\'),
                    indent(4, fr'printf "{self.hotstart.binary} -i {self.hotstart.iteration}\n" >>' + '${SLURM_JOB_FILE};' + '\\'),
                    indent(4, r'printf "popd\n" >>' + '${SLURM_JOB_FILE};' + '\\'),
                    indent(4, fr'printf "mv {self._relpath / self.hotstart.path.name if hasattr(self, "_relpath") else self.hotstart.path.resolve()} ./hotstart.nc\n" >> ' + '${SLURM_JOB_FILE};' + '\\'),
                ]
            return f

        return '\n'.join([
            'slurm:',
            indent(4, '@set -e;\\'),
            indent(4, r'printf "#!/bin/bash --login\n" > ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'printf "#SBATCH -D .\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'if [ ! -z "${SLURM_ACCOUNT}" ];'+'\\'),
            indent(4, 'then ' + '\\'),
            indent(8, r'printf "#SBATCH -A ${SLURM_ACCOUNT}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'fi;' + '\\'),
            indent(4, r'if [ ! -z "${SLURM_MAIL_USER}" ];'+'\\'),
            indent(4, 'then ' + '\\'),
            indent(8, r'printf "#SBATCH --mail-user=${SLURM_MAIL_USER}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(8, r'printf "#SBATCH --mail-type=${SLURM_MAIL_TYPE:-all}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'fi;' + '\\'),
            indent(4, r'printf "#SBATCH --output=${SLURM_LOG_FILE}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'printf "#SBATCH -n ${SLURM_NTASKS}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'if [ ! -z "${SLURM_WALLTIME}" ];'+'\\'),
            indent(4, 'then ' + '\\'),
            indent(8, r'printf "#SBATCH --time=${SLURM_WALLTIME}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'fi;' + '\\'),
            indent(4, r'if [ ! -z "${SLURM_PARTITION}" ];'+'\\'),
            indent(4, 'then ' + '\\'),
            indent(8, r'printf "#SBATCH --partition=${SLURM_PARTITION}\n" >> ${SLURM_JOB_FILE};' + '\\'),
            indent(4, r'fi;' + '\\'),
            indent(4, r'printf "\nset -e\n" >> ${SLURM_JOB_FILE};' + '\\'),
            *modules(),
            *hotstart(),
            indent(4, r'printf "${MPI_LAUNCHER} ${SCHISM_BINARY}" >> ${SLURM_JOB_FILE}')
        ])

    @property
    def run(self):
        f = [
            '',
            'run: slurm',
            '    @set -e;\\',
        ]

        return '\n'.join([line.replace("    ", "\t") for line in f]) + r"""
    touch ${SLURM_LOG_FILE};\
    rm -rf outputs/mirror.out outputs/fatal.error;\
    touch outputs/mirror.out outputs/fatal.error;\
    eval 'tail -f ${SLURM_LOG_FILE} outputs/mirror.out outputs/fatal.error &';\
    tail_pid=$${!};\
    job_id=$$(sbatch ${SLURM_JOB_FILE});\
    printf "$${job_id}\n";\
    job_id=$$(echo $${job_id} | awk '{print $$NF}');\
    ctrl_c() { \
        scancel "$${job_id}";\
    };\
    while [ $$(squeue -j $${job_id} | wc -l) -eq 2 ];\
    do \
        trap ctrl_c SIGINT;\
    done;\
    kill "$${tail_pid}"
"""


class MakefileDriver:

    def __new__(cls, server_config: ServerConfig = None, hotstart=None):
        if isinstance(server_config, SlurmConfig):
            return SlurmMakefile(server_config, hotstart=hotstart)
        return DefaultMakefile(server_config, hotstart=hotstart)
